import time

import gradio as gr
import numpy as np
from PIL import Image
from ocr_process.operators import ClasResizeImage
from ocr_process.utils.img_utils import np2base64
from openai import OpenAI

from taskinfo.uv_pipeline import UvPipeline

uv_predict = UvPipeline(model_path=r"D:\ocr\onnx\deploy\uvdoc\inference.onnx", use_openvino=True)

resize_func = ClasResizeImage(960)
def predict(image: Image.Image, prompt: str):
    image_data = np.array(image)
    result = uv_predict(image_data)
    if result['output_imgs']:
        image_data = result['output_imgs'][0]

    # img = Image.fromarray(image_data)
    # r_channel = img.split()[0]
    # rgb_image = r_channel.convert('RGB')
    # image_data = np.array(rgb_image)
    start_time = time.perf_counter()
    print(f"before shap={image_data.shape}")
    image_data = resize_func({"image": image_data})["image"]
    print(f"after shap={image_data.shape}")
    text_tesult = inference_with_api(np2base64(image_data), prompt=prompt)
    print(f"耗时：{time.perf_counter() - start_time}")
    return image_data, text_tesult


def inference_with_api(base64_image, prompt, sys_prompt="你是一个文档信息提取助手，提取文档中信息来回答问题。",
                       model_id="Qwen2.5-VL-7B-Instruct-GPTQ-Int4"):
    client = OpenAI(
        api_key="test_sk",
        base_url="http://39.98.50.56:8020/v1",
    )
    messages = [
        {"role": "system", "content": sys_prompt},
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                },
                {"type": "text", "text": prompt},
        #         报关单信息提取，提取多个报关商品信息，信息中包含‘项号，商品编号，商品名称及规格型号，数量及单位，单价 / 总价 / 币制，原产国(地区)，最终目的国(
        # 地区)，境内目的地，征免’，请你记住你的输出是一个json格式的字符串
                # {"type": "text", "text": f"文档信息提取,提取内容只能包含'{prompt}',请你记住你的输出是一个json格式的字符串"},
                # {"type": "text", "text": f"文档信息提取,提取内容只能包含'{prompt}'，时间转换成格式为yyyy-MM-dd HH:mm:ss,请你记住你的输出是一个json格式的字符串"},
            ],
        }
    ]
    completion = client.chat.completions.create(
        model=model_id,
        messages=messages,
        # temperature=0.7,

    )
    print(completion)
    content: str = completion.choices[0].message.content
    # try:
    #     bot_response = content[content.find("{"):]
    #     bot_response = bot_response[: bot_response.find("}") + 1]
    #     return bot_response
    # except:
    #     pass
    return content


def gradio_interface(image, prompt):
    return predict(image, prompt)


iface = gr.Interface(
    fn=gradio_interface,
    inputs=["image", gr.Textbox(label="我是一个报关单审核员，有什么问题可以问我", value="车号,净重", placeholder="提取文本以,分割")],
    outputs=['image', 'text'],
    title="报关单信息提取",
    description="上传一张图片，提取关键信息"
)

if __name__ == '__main__':
    iface.launch(server_name="0.0.0.0")
